import numpy as np
import cv2
from sklearn.model_selection import train_test_split
from sklearn.svm import SVC
from sklearn.metrics import accuracy_score
import pickle

# 读取图像
image = cv2.imread('Sandstone_1.tif', cv2.IMREAD_GRAYSCALE)
label_image = cv2.imread('Sandstone_1_segment.tif', cv2.IMREAD_GRAYSCALE)

print(f"图像形状：{image.shape}")

# 应用滤波器
def apply_filters(image):
    filtered_images = []
    # 均值滤波
    mean_filtered = cv2.blur(image, (5, 5))
    filtered_images.append(mean_filtered)
    
    # 高斯滤波
    gaussian_filtered = cv2.GaussianBlur(image, (5, 5), 0)
    filtered_images.append(gaussian_filtered)
    
    # Sobel滤波
    sobelx = cv2.Sobel(image, cv2.CV_64F, 1, 0, ksize=5)
    sobely = cv2.Sobel(image, cv2.CV_64F, 0, 1, ksize=5)
    sobel_filtered = np.sqrt(sobelx**2 + sobely**2)
    filtered_images.append(sobel_filtered)
    
    # Canny滤波
    canny_filtered = cv2.Canny(image, 100, 200)
    filtered_images.append(canny_filtered)
    
    return np.stack(filtered_images, axis=-1)

# 提取特征和标签
features = []
labels = []
filtered_images = apply_filters(image)

for i in range(image.shape[0]):
    for j in range(image.shape[1]):
        # 提取特征
        pixel_features = filtered_images[i, j, :]
        features.append(pixel_features)
        
        # 获取标签
        label = label_image[i, j]
        labels.append(label)

# 转换为NumPy数组
X = np.array(features)
y = np.array(labels)
print("完成从砂岩截面图1及其对应分区中获取X和y")

# 划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
print("完成train_test_split")

# 训练模型
clf = SVC(gamma='scale')
clf.fit(X_train, y_train)
print("完成随机SVM的训练")
# 预测测试集
y_pred = clf.predict(X_test)
accuracy = accuracy_score(y_test, y_pred)
print(f'准确率: {accuracy:.4f}')
# 保存模型
with open('clf.pickle', 'wb') as f:
    pickle.dump(clf, f)
print("已保存保存SVM模型到硬盘")

